#include <cassert>
#include <iostream>

#include "../src/kl.hh"
#include "../src/ordinallogistic.hh"

namespace 
{
  using namespace playerpiano;

  struct LogPTestData
  {
    struct
      {
        float w[17 + 5];
        Feature f[4];
      }                 inputs;
  
    float logp[6];
  };
  
  void
  test_logp (void)
  {
    LogPTestData data[] = 
      {{{{0.0312916, 0.200644, 0.567809, 0.200262, 0.21638, 0.585703,  0.66993, 0.139278, 0.539136, 0.652288, 0.124868, 0.626049,  0.127142, 0.038082, 0.993499, 0.209523, 0.352901, 0.365108,  0.0927866, 0.314479, 0.92601,  0.544863}, {{6, 0.162308}, {5, 0.0746401}, {7, 0.894719}, {13,  0.344219}}}, {-1.89373, -1.66839, -1.61241, -1.92359, -1.86722, -1.82698}}, {{{0.594499, 0.874378, 0.678339, 0.758516, 0.924569,  0.735099, 0.139203, 0.106227, 0.799701, 0.10905, 0.012061,  0.0681453, 0.806202, 0.899527, 0.65916, 0.703037, 0.713416,  0.585049, 0.733149, 0.158174, 0.551108,  0.510409}, {{10, 0.838431}, {11, 0.813955}, {11, 0.956609}, {11,  0.636031}}}, {-7.97756, -6.26571, -4.83374, -3.20573, -1.65789, -0.276109}}, {{{0.160092, 0.0554398, 0.0320392, 0.900932, 0.0208889,  0.949213, 0.232338, 0.791881, 0.00882791, 0.881067, 0.426136,  0.892354, 0.349668, 0.17803, 0.71272, 0.307305, 0.616519,  0.0198555, 0.161612, 0.796897, 0.778088,  0.2059}, {{4, 0.205004}, {1, 0.160866}, {16, 0.617997}, {12,  0.15046}}}, {-2.20739, -2.01925, -1.72646, -1.41027, -1.96297, -1.63563}}, {{{0.172964, 0.259934, 0.597108, 0.201248, 0.940626,  0.468053, 0.58828, 0.320181, 0.514491, 0.575699, 0.238612,  0.142151, 0.801771, 0.268393, 0.622093, 0.122295, 0.640158,  0.471497, 0.844005, 0.916395, 0.435155,  0.310631}, {{6, 0.226008}, {4, 0.765935}, {8, 0.26219}, {6,  0.0506967}}}, {-2.37049, -1.9928, -1.70208, -1.74854, -1.39913, -1.78911}}, {{{0.6289, 0.564687, 0.321564, 0.582644, 0.0406205,  0.244506, 0.807074, 0.00694498, 0.802009, 0.102356, 0.00530305,  0.738552, 0.179916, 0.98006, 0.365145, 0.267055, 0.335911,  0.0636651, 0.92999, 0.956424, 0.109903,  0.29773}, {{12, 0.667799}, {1, 0.905727}, {16, 0.481003}, {0,  0.733043}}}, {-3.89767, -3.43062, -2.89936, -2.12498, -1.61168, -0.55635}}, {{{0.346235, 0.323083, 0.440383, 0.488537, 0.539162,  0.316138, 0.638374, 0.386181, 0.533859, 0.577587, 0.458458,  0.406121, 0.168714, 0.310532, 0.122546, 0.342456, 0.238724,  0.354108, 0.0126428, 0.0447257, 0.570924,  0.448381}, {{11, 0.53164}, {2, 0.311682}, {7, 0.224689}, {0,  0.125298}}}, {-1.6434, -1.66483, -1.66312, -1.7787, -1.94243, -2.15679}}, {{{0.0912572, 0.823145, 0.685528, 0.809159, 0.452884,  0.436964, 0.151669, 0.231572, 0.994426, 0.0308431, 0.982955,  0.92104, 0.87188, 0.688387, 0.744231, 0.566932, 0.859237,  0.643662, 0.173307, 0.118551, 0.327597,  0.331979}, {{4, 0.948618}, {7, 0.993253}, {5, 0.23634}, {1,  0.508834}}}, {-3.9268, -2.81331, -2.43172, -1.9125, -1.51692, -0.765062}}, {{{0.26309, 0.184094, 0.783456, 0.0718695, 0.111421,  0.952522, 0.78903, 0.0410264, 0.128466, 0.0314815, 0.917151,  0.352639, 0.384234, 0.464549, 0.0579138, 0.708978, 0.210927,  0.345998, 0.730317, 0.376998, 0.26231,  0.352745}, {{10, 0.493977}, {3, 0.868165}, {15, 0.99922}, {16,  0.168651}}}, {-3.30559, -2.78533, -2.18608, -2.1862, -1.47472, -0.8028}}, {{{0.710521, 0.796295, 0.887799, 0.216129, 0.92149,  0.755269, 0.759333, 0.184648, 0.00433954, 0.40263, 0.375099,  0.720098, 0.946426, 0.693652, 0.164172, 0.3741, 0.216109,  0.316654, 0.901862, 0.0213546, 0.722132,  0.448489}, {{10, 0.902642}, {6, 0.852704}, {8, 0.0116114}, {1,  0.652194}}}, {-4.3984, -3.65392, -2.99522, -2.42803, -1.18916, -0.655657}}, {{{0.0148435, 0.636575, 0.0901212, 0.896925, 0.25551,  0.451927, 0.0857817, 0.494295, 0.880411, 0.731829, 0.139356,  0.800643, 0.716239, 0.357729, 0.923247, 0.48399, 0.814377,  0.336375, 0.201115, 0.0355007, 0.911735,  0.483671}, {{10, 0.189504}, {16, 0.383307}, {6, 0.896891}, {14,  0.847097}}}, {-4.0109, -3.00047, -2.61177, -1.67662, -1.54827, -0.778503}}}; 

    for (unsigned int i = 0; i < sizeof (data) / sizeof (data[0]); ++i)
      {
        OrdinalLogisticRegression r (17, 6, 1, data[i].inputs.w + 5, data[i].inputs.w, 1);
        std::vector<Feature> features (data[i].inputs.f, data[i].inputs.f + 4);
        IndexedFeatureSet feature_set;
        unsigned int ns = 256 * drand48 ();
        feature_set[ns] = features;

        std::vector<float> logp = r.logp (feature_set);
  
        for (unsigned int j = 0; j < 5; ++j)
          {
            assert (fabs (logp[j] - data[i].logp[j]) <=
                    1e-4 * (1 + fabs (logp[j]) + fabs (data[i].logp[j])) ||
                    (std::cerr << "logp[" << j << "] " << logp[j] 
                               << " ?= " << data[i].logp[j] << std::endl,
                     0));
          }
      }
  }

  struct GradWTestData
  {
    struct
      {
        float        w[9 + 3];
        Feature      f[2];
        float        q[4];
        unsigned int nitems;
      }                 inputs;
  
    struct
      {
        float logp[4];
        float kl;
        float gradw[9 + 3];
      }                 outputs;
  };

  void
  test_gradw (void)
  {
    GradWTestData data[] = 
      {{{{0.410369, 0.898394, 0.0888508, 0.270436, 0.996413, 0.935319,  0.813065, 0.943119, 0.382258, 0.501914, 0.318747,  0.889009}, {{5, 0.913255}, {1, 0.663222}}, {0.431776, 0.204129,  0.103542, 0.260553},  3}, {{-2.29986, -1.70029, -1.58874, -0.667648}, -1.67856, {0.194713, 0.0535364, 0.222743, 0, -0.621349, 0, 0, 0, -0.855596, 0,  0, 0}}}, {{{0.573829, 0.51311, 0.569005, 0.0396147, 0.84967,  0.908834, 0.158636, 0.141221, 0.76082, 0.638398, 0.162224,  0.205902}, {{2, 0.947755}, {2, 0.695279}}, {0.267182, 0.241155,  0.21547, 0.276193},  6}, {{-3.29699, -2.37758, -1.39744, -0.473205}, -1.88607, {0.134549, 0.29305, 0.251976, 0, 0, -1.57003, 0, 0, 0, 0, 0,  0}}}, {{{0.86671, 0.0407655, 0.810119, 0.419127, 0.670337,  0.546611, 0.23629, 0.906016, 0.101332, 0.506996, 0.38662,  0.997183}, {{8, 0.942696}, {6, 0.365775}}, {0.325103, 0.186388,  0.405455, 0.0830541},  5}, {{-2.43196, -2.17318, -1.08846, -0.77309}, -1.70122, {0.063896,  0.301658, 0.216507, 0, 0, 0, 0, 0, 0, -0.338554,  0, -0.872541}}}, {{{0.678046, 0.663506, 0.000507453, 0.455886,  0.0490386, 0.857236, 0.133797, 0.41512, 0.23892, 0.43811, 0.46346,  0.868509}, {{8, 0.00262933}, {1, 0.532093}}, {0.19318, 0.192852,  0.328615, 0.285352},  6}, {{-0.734026, -1.38369, -2.01882, -1.99096}, -1.64019, {-0.3998, -0.455176, -0.148872, 0, 0.415124, 0, 0, 0, 0, 0, 0,  0.00205133}}}, {{{0.419432, 0.995737, 0.990209, 0.176125,  0.638959, 0.835864, 0.312163, 0.512619, 0.638452, 0.379979,  0.263124, 0.655383}, {{2, 0.504655}, {0, 0.964858}}, {0.063861,  0.573247, 0.108687, 0.254205},  3}, {{-1.25671, -1.08438, -1.48836, -1.88681}, -1.34327, {-0.36054, -0.317506, -0.43272, 0.298116, 0, 0.155925, 0, 0, 0, 0, 0,  0}}}, {{{0.0215754, 0.68518, 0.679067, 0.734836, 0.405566,  0.15027, 0.259635, 0.739099, 0.415357, 0.974144, 0.620675,  0.903235}, {{5, 0.103194}, {1, 0.461525}}, {0.311649, 0.166024,  0.266546, 0.255781},  4}, {{-1.25616, -1.04769, -1.50283, -1.95186}, -1.46525, {0.0215102, -0.329113, -0.283538, 0, 0.112928, 0, 0, 0, 0.02525, 0, 0,  0}}}, {{{0.477568, 0.558398, 0.815865, 0.588869, 0.436373,  0.462049, 0.79429, 0.903689, 0.757307, 0.727212, 0.388723,  0.753419}, {{0, 0.497672}, {4, 0.988113}}, {0.358578, 0.287077,  0.323077, 0.0312682},  3}, {{-2.51308, -1.80463, -1.17702, -0.806874}, -1.8247, {0.11837,  0.213969, 0.143027, -0.543778, 0, 0, 0, -1.07966, 0, 0, 0,  0}}}, {{{0.870172, 0.317749, 0.894774, 0.561622, 0.0301027,  0.511607, 0.417206, 0.00322357, 0.214238, 0.922738, 0.980832,  0.541175}, {{4, 0.419948}, {3, 0.0190489}}, {0.167517, 0.61001,  0.0234008, 0.199072},  2}, {{-0.621378, -1.48225, -1.7907, -2.67617}, -1.58293, {-0.804772, -0.145677, -0.577632, 0, 0, 0, 0.00927177, 0.204403, 0, 0,  0, 0}}}, {{{0.725853, 0.82585, 0.0578584, 0.486355, 0.848856,  0.740972, 0.187686, 0.168606, 0.954082, 0.17935, 0.157583,  0.656998}, {{2, 0.536876}, {8, 0.176126}}, {0.328852, 0.255965,  0.193838, 0.221345},  5}, {{-1.20411, -1.41644, -1.72877, -1.2731}, -1.37543, {-0.116277, -0.122887, 0.0470457, 0, 0, -0.0696835, 0, 0, 0, 0,  0, -0.0228602}}}, {{{0.523397, 0.715212, 0.332519, 0.820989,  0.492173, 0.449582, 0.606665, 0.995139, 0.434314, 0.963227,  0.757809, 0.254168}, {{2, 0.246628}, {7, 0.794621}}, {0.727215,  0.0676953, 0.0805681, 0.124522},  2}, {{-1.64477, -1.45512, -1.45728, -1.07675}, -1.5461, {0.27246,  0.0108749, 0.0499206, 0, 0, -0.275932, 0, 0, 0, 0, -0.889038,  0}}}}; 

    for (unsigned int i = 0; i < sizeof (data) / sizeof (data[0]); ++i)
      {
        OrdinalLogisticRegression rzero (9, 4, data[i].inputs.nitems, data[i].inputs.w + 3, data[i].inputs.w, 1);
        std::vector<Feature> features (data[i].inputs.f, data[i].inputs.f + 2);
        IndexedFeatureSet feature_set;
        unsigned int ns = 256 * drand48 ();
        feature_set[ns] = features;
        std::vector<float> logp = rzero.logp (feature_set);

        std::vector<float> w (data[i].inputs.w, data[i].inputs.w + 9 + 3);
        OrdinalLogisticRegression r (9, 4, data[i].inputs.nitems, &w[3], &w[0], 1);
        std::vector<float> q (data[i].inputs.q, data[i].inputs.q + 4);

        float kl = kldivergence (logp, q);
          
        r.update (1, 1, 0, feature_set, logp, q);

        assert (fabs (kl - data[i].outputs.kl) <=
                1e-4 * (1 + fabs (kl) + fabs (data[i].outputs.kl)) ||
                (std::cerr << "kl = " << kl
                           << " ?= " << data[i].outputs.kl << std::endl,
                 0));

        for (unsigned int j = 0; j < 4; ++j)
          {
            assert (fabs (logp[j] - data[i].outputs.logp[j]) <=
                    1e-4 * (1 + fabs (logp[j]) + fabs (data[i].outputs.logp[j])) ||
                    (std::cerr << "logp[" << j << "] " << logp[j] 
                               << " ?= " << data[i].outputs.logp[j] << std::endl,
                     0));
          }

        float xdotx = 0;

        for (unsigned int k = 0; k < 2; ++k)
          { 
            xdotx += features[k].value;
          }

        OrdinalLogisticRegression::DFunc dfunc (logp,
                                                q,
                                                data[i].inputs.w,
                                                xdotx,
                                                data[i].inputs.nitems,
                                                1);
        PianoVec s0 (4, 0);
        PianoVec sprime0 (dfunc (0, s0));

        for (unsigned int j = 0; j < 3; ++j)
          {
            assert (fabs (sprime0[j + 1] - data[i].outputs.gradw[j]) <=
                    1e-4 * (1 + fabs (sprime0[j + 1]) + fabs (data[i].outputs.gradw[j])) ||
                    (std::cerr << "sprime0[" << j + 1 << "] " << sprime0[j + 1] 
                               << " ?= " << data[i].outputs.gradw[j] << std::endl,
                     0));
          }

        std::vector<float> dw (9, 0);

        for (unsigned int j = 0; j < 2; ++j)
          {
            dw[features[j].feature] += sprime0[0] * features[j].value;
          }

        for (unsigned int j = 0; j < 9; ++j)
          {
            assert (fabs (dw[j] - data[i].outputs.gradw[j + 3]) <=
                    1e-4 * (1 + fabs (dw[j]) + fabs (data[i].outputs.gradw[j + 3])) ||
                    (std::cerr << "dw[" << j << "] " << dw[j] 
                               << " ?= " << data[i].outputs.gradw[j + 3] << std::endl,
                     0));
          }
      }
  }

  struct UpdateTestData
  {
    struct
      {
        float        w[9 + 3];
        Feature      f[2];
        float        q[4];
        unsigned int nitems;
      }                 inputs;
  
    struct
      {
        float wnew[9 + 3];
      }                 outputs;
  };

  void
  test_update (void)
  {
    UpdateTestData data[] = 
      {{{{0.449187, 0.480721, 0.210087, 0.118627, 0.360142, 0.343098,  0.943237, 0.219936, 0.214443, 0.939736, 0.23243,  0.956196}, {{8, 0.592141}, {3, 0.251585}}, {0.463549, 0.27865,  0.250737, 0.00706395},  2}, {{0.472831, 0.530597, 0.373309, 0.118627, 0.360142, 0.343098,  0.718417, 0.219936, 0.214443, 0.939736, 0.23243,  0.427051}}}, {{{0.08252, 0.0658325, 0.766196, 0.737674, 0.0687064,  0.872388, 0.317009, 0.256953, 0.858619, 0.753761, 0.956867,  0.913855}, {{7, 0.915383}, {0, 0.533826}}, {0.249376, 0.327201,  0.2294, 0.194023},  1}, {{0.123505, 0.209172, 0.410666, 0.367511, 0.0687064, 0.872388,  0.317009, 0.256953, 0.858619, 0.753761, 0.322126,  0.913855}}}, {{{0.150284, 0.722534, 0.82881, 0.0641854, 0.688272,  0.709518, 0.74629, 0.998353, 0.922076, 0.971844, 0.677584,  0.125965}, {{7, 0.605067}, {1, 0.714891}}, {0.310166, 0.140965,  0.245493, 0.303376},  6}, {{0.277406, 0.651438, 0.699648, 0.0641854, 0.540869, 0.709518,  0.74629, 0.998353, 0.922076, 0.971844, 0.552825,  0.125965}}}, {{{0.903581, 0.838378, 0.905775, 0.826917, 0.220629,  0.260748, 0.755491, 0.104383, 0.391819, 0.196563, 0.0672189,  0.394865}, {{7, 0.645529}, {4, 0.19821}}, {0.0902431, 0.263015,  0.601823, 0.0449185},  2}, {{0.31968, 0.273991, 0.620049, 0.826917, 0.220629, 0.260748,  0.755491, 0.209529, 0.391819, 0.196563, 0.409659,  0.394865}}}, {{{0.540076, 0.70813, 0.148981, 0.700041, 0.891877,  0.907094, 0.245399, 0.861663, 0.986102, 0.0801775, 0.0247707,  0.600915}, {{0, 0.23061}, {2, 0.975795}}, {0.355267, 0.226957,  0.0917093, 0.326067},  6}, {{0.616528, 0.77169, 0.185426, 0.600274, 0.891877, 0.484945,  0.245399, 0.861663, 0.986102, 0.0801775, 0.0247707,  0.600915}}}, {{{0.987424, 0.206142, 0.0182486, 0.157908,  0.0194787, 0.133897, 0.478172, 0.449778, 0.870498, 0.433856,  0.586296, 0.542683}, {{4, 0.625099}, {6, 0.572193}}, {0.227838,  0.17557, 0.227888, 0.368704},  4}, {{0.768297, 0.171236, 0.023988, 0.157908, 0.0194787, 0.133897,  0.478172, 0.436846, 0.870498, 0.422019, 0.586296,  0.542683}}}, {{{0.369584, 0.486711, 0.967376, 0.566927, 0.206192,  0.905782, 0.979952, 0.360784, 0.187944, 0.747874, 0.960474,  0.226887}, {{6, 0.709771}, {6, 0.298096}}, {0.0510675, 0.450101,  0.0700809, 0.42875},  1}, {{0.0952961, 0.29479, 0.324848, 0.566927, 0.206192, 0.905782,  0.979952, 0.360784, 0.187944, 0.70317, 0.960474,  0.226887}}}, {{{0.464877, 0.220838, 0.523281, 0.292907, 0.864549,  0.24956, 0.153697, 0.806195, 0.897173, 0.682633, 0.947505,  0.900413}, {{0, 0.917221}, {8, 0.321849}}, {0.38678, 0.0776753,  0.487189, 0.0483558},  4}, {{0.478181, 0.152872, 0.573075, 0.0145275, 0.864549, 0.24956,  0.153697, 0.806195, 0.897173, 0.682633, 0.947505,  0.802731}}}, {{{0.0497906, 0.854443, 0.866772, 0.30193, 0.926315,  0.0990308, 0.401895, 0.081092, 0.403034, 0.806124, 0.537346,  0.831532}, {{2, 0.249337}, {7, 0.999929}}, {0.537771, 0.125081,  0.253551, 0.083597},  3}, {{0.24676, 0.641953, 0.67508, 0.30193, 0.926315, 0.0104604,  0.401895, 0.081092, 0.403034, 0.806124, 0.182147,  0.831532}}}, {{{0.722952, 0.82705, 0.54227, 0.946976, 0.766204,  0.732088, 0.492479, 0.0925325, 0.899433, 0.430158, 0.566164,  0.993502}, {{6, 0.497538}, {2, 0.349066}}, {0.089228, 0.102491,  0.525202, 0.283079},  2}, {{0.336498, 0.297096, 0.340541, 0.946976, 0.766204, 0.872205,  0.492479, 0.0925325, 0.899433, 0.629872, 0.566164, 0.993502}}}}; 

    for (unsigned int i = 0; i < sizeof (data) / sizeof (data[0]); ++i)
      {
        std::vector<float> w (data[i].inputs.w, data[i].inputs.w + 9 + 3);
        OrdinalLogisticRegression rzero (9, 4, data[i].inputs.nitems, &w[3], &w[0], 1);
        std::vector<Feature> features (data[i].inputs.f, data[i].inputs.f + 2);
        IndexedFeatureSet feature_set;
        unsigned int ns = 256 * drand48 ();
        feature_set[ns] = features;
        std::vector<float> logp = rzero.logp (feature_set);
        std::vector<float> q (data[i].inputs.q, data[i].inputs.q + 4);

        rzero.update (1, 1, 0, feature_set, logp, q);
        assert (rzero.failures () == 0);
  
        for (unsigned int j = 0; j < 9 + 3; ++j)
          {
            assert (fabs (w[j] - data[i].outputs.wnew[j]) <=
                    6e-3 * (1.0 + fabs (w[j]) + fabs (data[i].outputs.wnew[j])) ||
                    (std::cerr << "(" << i << ") w[" << j << "] " << w[j] 
                               << " ?= " << data[i].outputs.wnew[j] << std::endl,
                     0));
          }
      }
  }
}

int
main (void)
{
  test_logp ();
  test_gradw ();
  test_update ();

  return 0;
}
